"""
手写数字识别 - 带数据增强的最优 KNN 模型选择器
通过数据增强扩充训练集，提升模型泛化能力
"""
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
import pickle
from tqdm import tqdm
import cv2

# 数据增强函数：对图像进行旋转、缩放、平移等变换
def augment_image(img, angle_range=(-10, 10), scale_range=(0.8, 1.2), shift_range=(-1, 1)):
    """
    对 8x8 手写数字图像进行增强
    :param img: 原始图像（1D 数组，长度 64）
    :param angle_range: 旋转角度范围
    :param scale_range: 缩放比例范围
    :param shift_range: 平移像素范围
    :return: 增强后的图像（1D 数组）
    """
    img_2d = img.reshape(8, 8)  # 转为 2D 图像
    augmented_imgs = []

    # 生成多个增强版本
    for _ in range(5):  # 每个图像生成 5 个增强样本
        # 随机旋转
        angle = np.random.uniform(angle_range[0], angle_range[1])
        M_rot = cv2.getRotationMatrix2D((4, 4), angle, 1.0)
        rotated = cv2.warpAffine(
            img_2d, M_rot, (8, 8),
            flags=cv2.INTER_NEAREST,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=16  # 背景设为 16（与原始数据一致）
        )

        # 随机缩放
        scale = np.random.uniform(scale_range[0], scale_range[1])
        new_size = (int(8 * scale), int(8 * scale))
        resized = cv2.resize(rotated, new_size, interpolation=cv2.INTER_NEAREST)
        
        # 随机平移
        shift_x = np.random.randint(shift_range[0], shift_range[1] + 1)
        shift_y = np.random.randint(shift_range[0], shift_range[1] + 1)
        M_trans = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
        shifted = cv2.warpAffine(
            resized, M_trans, (8, 8),
            flags=cv2.INTER_NEAREST,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=16
        )

        # 裁剪或填充到 8x8
        padded = np.ones((8, 8)) * 16
        h, w = shifted.shape
        y_start = max(0, (8 - h) // 2)
        x_start = max(0, (8 - w) // 2)
        y_end = min(8, y_start + h)
        x_end = min(8, x_start + w)
        padded[y_start:y_end, x_start:x_end] = shifted[:y_end-y_start, :x_end-x_start]

        augmented_imgs.append(padded.flatten())
    return augmented_imgs

# 加载数据集
print("加载手写数字数据集...")
digits = load_digits()
X = digits.data
y = digits.target

# 划分训练集和测试集
print("划分训练集和测试集...")
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# 数据增强：扩充训练集
print("数据增强中...")
X_train_augmented = []
y_train_augmented = []
for img, label in zip(X_train, y_train):
    X_train_augmented.append(img)
    y_train_augmented.append(label)
    # 对每个图像生成增强样本
    augmented_samples = augment_image(img)
    X_train_augmented.extend(augmented_samples)
    y_train_augmented.extend([label] * len(augmented_samples))

X_train_augmented = np.array(X_train_augmented)
y_train_augmented = np.array(y_train_augmented)
print(f"增强后训练集大小: {X_train_augmented.shape[0]}")

# 数据标准化
print("数据标准化...")
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train_augmented)
X_test_scaled = scaler.transform(X_test)

# 寻找最优 K 值
print("寻找最优 K 值...")
best_accuracy = 0
best_k = 0
best_model = None

for k in tqdm(range(1, 21), desc="K 值搜索"):
    knn = KNeighborsClassifier(n_neighbors=k, metric='manhattan', n_jobs=-1)
    knn.fit(X_train_scaled, y_train_augmented)
    y_pred = knn.predict(X_test_scaled)
    accuracy = accuracy_score(y_test, y_pred)
    if accuracy > best_accuracy:
        best_accuracy = accuracy
        best_k = k
        best_model = knn

print(f"最优 K 值: {best_k}, 准确率: {best_accuracy:.4f}")

# 保存模型和标准化器
print("保存模型和标准化器...")
with open('best_knn_model.pkl', 'wb') as f:
    pickle.dump(best_model, f)
with open('scaler.pkl', 'wb') as f:
    pickle.dump(scaler, f)
print("模型和标准化器已保存")